package ink.zhaohui.javaStudy.instrument;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.*;

import ink.zhaohui.javaStudy.annotation.Monitor;
import javassist.*;

public class MonitorTransformer implements ClassFileTransformer {
  private List<String> monitorClasses;

  private final String[] WHITELIST_PREFIXES = new String[] {};

  private final String[] BLACKLIST_PREFIXES = new String[] {"jdk/", "java/", "javax/", "com/sun/","sun/",
      "org/apache/maven/", "junit/", "org/junit/", "org/testng/"};

  public void setMonitorClasses(List<String> monitorClasses) {
    this.monitorClasses = monitorClasses;
  }

  public List<String> getMonitorClasses() {
    return monitorClasses;
  }

  @Override
  public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
      ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
    if (isSkipClass(className)) {
      // Ignore system class and reloaded class
      return null;
    }
    //
    ClassPool pool = ClassPool.getDefault();
    String cn = className.replace("/", ".");
    System.out.println("className["+className+"]cn["+cn+"].");
    byte[] transform = null;
    CtClass cc = null;
    try {
      cc = pool.get(cn);
      transform = addMonitor(cc);
    } catch (Exception e) {
      // TODO Auto-generated catch block
      e.printStackTrace();
    } finally {
      if (cc != null) {
        //cc.detach();
      }
    }
    //

    return transform;
  }

  private byte[] addMonitor(CtClass cc) throws Exception {
    //
    CtMethod[] cm = null;
    cm = cc.getDeclaredMethods();
    for (CtMethod m : cm) {
      Monitor a = (Monitor) m.getAnnotation(Monitor.class);
      System.out.println("class["+cc.getName()+"]Method["+m.getName()+"].");
      if (a != null) {
        Monitor.Type[] ts = a.see();
        if (hasType(ts, Monitor.Type.PARAMS)) {
          
          m.insertBefore("{ System.out.print(\"params[\"); "+
                          " for(int i=0 ; i < $args.length ; i++ ) { "+
                          "     System.out.print($args[i].toString()+\"|\"); }"+
                          " System.out.println(\"]\"); }");
        }
        if (hasType(ts, Monitor.Type.SPEED)) {
          //long _t1_=System.currentTimeMillis();
          //long _t2_=System.currentTimeMillis();
          //System.out.println("["+m.getName()+"]Millis["+(_t2_-_t1_)+"]");

          m.addLocalVariable("_t1_",CtClass.longType);
          
          m.insertBefore("{ _t1_=System.currentTimeMillis(); }");

          m.insertAfter("{ _t1_=System.currentTimeMillis()-_t1_;  "
                         +"System.out.println(\"["+cc.getName()+"|"+m.getName()+"]Millis[\"+_t1_+\"]\");  }");
          
        }
      }
    }
    //
    return cc.toBytecode();
  }

  private boolean hasType(Monitor.Type[] ts, Monitor.Type t) {
    for (Monitor.Type t1 : ts) {
      if (t1 == t) {
        return true;
      }
    }
    return false;
  }

  private boolean isSkipClass(String className) {
    // className can be null for Java 8 lambdas
    if (null == className) {
      return true;
    }
    for (String prefix : WHITELIST_PREFIXES) {
      if (className.startsWith(prefix)) {
        return true;
      }
    }
    for (String prefix : BLACKLIST_PREFIXES) {
      if (className.startsWith(prefix)) {
        return true;
      }
    }
    return false;
  }
}
